"""
工具函数模块
"""
import random
import numpy as np
import torch
import torch.backends.cudnn

def set_seed(seed):
    """
    设置随机种子以确保结果可复现
    
    Args:
        seed (int): 随机种子值
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def print_model_info(model):
    """
    打印模型信息
    
    Args:
        model: PyTorch模型
    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"模型总参数数量: {total_params:,}")
    print(f"可训练参数数量: {trainable_params:,}")
    print(f"模型大小: {total_params * 4 / 1024 / 1024:.2f} MB")

def format_time(seconds):
    """
    格式化时间显示
    
    Args:
        seconds (float): 秒数
        
    Returns:
        str: 格式化的时间字符串
    """
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = int(seconds % 60)
    
    if hours > 0:
        return f"{hours}h {minutes}m {seconds}s"
    elif minutes > 0:
        return f"{minutes}m {seconds}s"
    else:
        return f"{seconds}s"

def create_train_test_split(total_batches, test_batches, max_train_batches=None):
    """
    创建训练和测试数据的批次索引
    
    Args:
        total_batches (int): 总批次数
        test_batches (int): 测试批次数
        max_train_batches (int, optional): 最大训练批次数
        
    Returns:
        tuple: (train_batch_indices, test_batch_indices)
    """
    all_batch_indices = list(range(total_batches))
    test_batch_indices = np.random.choice(all_batch_indices, test_batches, replace=False)
    
    train_batch_indices = [idx for idx in all_batch_indices if idx not in test_batch_indices]
    
    if max_train_batches is not None:
        train_batch_indices = [idx for idx in train_batch_indices if idx < max_train_batches]
    
    return train_batch_indices, test_batch_indices 